import os
import numpy as np
import pandas as pd
from loguru import logger

from models.global_config import glm
from models.stock_model import StockNumber, DayInfo


def get_all_stock(N, stock):
    df = pd.read_csv('all.csv')
    count = 0
    for row in df.index:
        if stock and df.loc[row]['ts_code'] != stock:
            continue
        sn = StockNumber(df.loc[row])
        if '退' in sn.name:
            continue
        if 'ST' in sn.name:
            continue
        if not str(sn.ts_code).startswith('60') \
                and not str(sn.ts_code).startswith('30') \
                and not str(sn.ts_code).startswith('00'):
            continue
        count += 1
        logger.info(f'{sn.ts_code},{sn.name}')
        analysis_stock(N, sn)


def analysis_stock(N, sn):
    csv_path = f'stocks/{sn.ts_code}.csv'
    df2 = pd.read_csv(csv_path)
    q10 = []
    q11 = []
    all_di = []
    parr = 0
    today_di = None
    cnt = 0
    pct_arr = []
    av_high = []
    av_low = []
    date_arr = []
    price_arr = []
    for row2 in df2.index:
        di = DayInfo(sn, df2.loc[row2])
        di.name = sn.name
        di.ts_code = sn.ts_code
        di.industry = sn.industry
        # di.trade_date = df2.loc[row2]['trade_date']
        # di.open = df2.loc[row2]['open']
        # di.high = df2.loc[row2]['high']
        # di.low = df2.loc[row2]['low']
        # di.close = float(df2.loc[row2]['close'])
        if di.close < 5:
            break
        if cnt > 40:
            data_dict = {
                'all_pct': pct_arr,
                'av_pct_high': [round(np.mean(av_high[2:-2]), 2)] * len(pct_arr),
                'av_pct_low': [round(np.mean(av_low[2:-2]), 2)] * len(pct_arr),
                'price': price_arr[:-1],
                'date': date_arr[:-1],
            }
            df = pd.DataFrame(data_dict)
            # 将数据写入CSV文件
            file_path = f'{sn.ts_code}_{sn.name}_{N}_{today_di.trade_date}.csv'
            if os.path.exists(file_path):
                os.remove(file_path)
            df.to_csv(file_path, index=False)
            break
        # di.pre_close = df2.loc[row2]['pre_close']
        # di.change = df2.loc[row2]['change']
        # di.pct_chg = float(df2.loc[row2]['pct_chg'])
        # di.vol = df2.loc[row2]['vol']
        # di.amount = df2.loc[row2]['amount']
        if len(q10) == N * 2 + 1:
            ad = all_di[N]
            if max(q10) == q10[N]:
                pct = round((parr - q10[N]) / q10[N] * 100, 2)
                pct_arr.append(pct)
                av_low.append(pct)
                price_arr.append(q10[N])
                date_arr.append(ad.trade_date)
                parr = q10[N]
                cnt += 1
            elif min(q11) == q11[N]:
                pct = round((parr - q11[N]) / q11[N] * 100, 2)
                pct_arr.append(pct)
                av_high.append(pct)
                price_arr.append(q11[N])
                date_arr.append(ad.trade_date)
                parr = q11[N]
                cnt += 1
            all_di.pop(0)
            q10.pop(0)
            q11.pop(0)
        all_di.append(di)
        q10.append(di.high)
        q11.append(di.low)
        if today_di is None:
            parr = di.low
            today_di = di
            price_arr.append(di.low)
            date_arr.append(di.trade_date)


def printNumber(stocks, N):
    for stock in stocks:
        get_all_stock(N, stock)


def run(bkd=0, stocks=None, N=6):
    stocks = [''] if stocks is None else stocks
    log_dir = os.path.abspath(__file__).replace('.py', '')
    if not os.path.exists(log_dir):
        os.mkdir(log_dir)

    logger.info('生成csv start')
    for idx, trade_date in enumerate(glm.all_trade_date_arr):
        if idx > bkd:
            break
        try:
            printNumber(stocks, N)
        except Exception as e:
            logger.error(e)
    logger.info('生成excel运行完成')
